import numpy as np
import matplotlib
matplotlib.use('TkAgg')#TkAgg是图形后端
from matplotlib import pyplot as plt
import random
#用于加载和保存数据，处理压缩文件和对象的序列化反序列化
import gzip
import _pickle
import pickle

class Activation:
    """
    激活函数类, 包括ReLU和Softmax激活函数, 其他函数暂未实现
    """
    #ReLU函数的核心特性：负值截断为零，正值不变。有助于解决梯度消失问题。
    @classmethod #@classmethod 是一个类方法的装饰器，表示这个方法是属于类（而不是实例）的方法。
    def ReLU(cls, Z):
        A = np.zeros_like(Z)#创建一个与Z具有相同形状和类型的零数组A。np.zeros_like(Z)会生成一个和Z一样形状的数组，并且数组中的每个元素都初始化为零。
        return np.maximum(Z, A)#np.maximum(Z, A) 保证输出大于0的值保持不变，其他值被设置为0。

    #Sigmoid激活函数
    @classmethod
    def Sigmoid(cls, Z):
        return 1 / (1 + 1 / np.exp(Z))
    #np.exp(Z) 会逐元素地将 Z 的每个元素计算成 e 的指数。
    @classmethod
    def Softmax(cls, Z):
        Z = np.exp(Z)  # 对输入Z进行指数运算
        Z_sum = Z.sum(axis=1).reshape((-1, 1))  # 对每一行求和
        return Z / Z_sum  # 将每一行的每个元素除以该行的总和，得到每个元素对应的概率


class Delta:
    """
    梯度求导类
    """

    def __init__(self, func):
        """
        类的初始化方法
        :param func: Activation类的函数指针
        它是 Activation 类中的一个方法（比如 ReLU 或 Sigmoid）。这个函数将用于计算激活函数的梯度（导数）。
        比如，如果传入的是 Activation.ReLU，那么 self.func 就会存储 Activation.ReLU，后续可以通过 self.func 来调用它。
        """
        self.func = func


    def delta(self, A):
        """
        这个方法计算给定激活函数的导数。
        A是输入的激活值（即从前一层传递过来的值，经过激活函数处理后的值）。
        对于 ReLU 激活函数，输出的梯度是一个和A形状相同的数组，其值为 1 或 0。
        对于已经通过 Sigmoid 函数的激活值A，它的导数（梯度）就是 A * (1 - A)，这就是 Sigmoid 函数的标准导数公式。
        """
        if self.func == Activation.ReLU:
            return (A > 0).astype(float) #会生成一个布尔值数组,True（即 A > 0）转换为 1.0，False（即 A <= 0）转换为 0.0
        elif self.func == Activation.Sigmoid:
            return A * (1 - A)


class Loss:
    """
    损失函数类
    Loss 类的作用是计算损失函数，并且提供一些与损失和准确度相关的功能。
    在神经网络中，损失函数用于衡量预测值与真实值之间的差距。损失越小，说明模型的预测越准确。
    """

    def __init__(self, y_hat, y):
        """
        :param y_hat: 预测值
        :param y: 标签值
        __init__(self, y_hat, y)：初始化方法，接收预测值 y_hat 和真实标签值 y。
        y_hat 是神经网络的输出（预测值），其形状通常是 (batch_size, num_classes)，表示每个样本在每个类别上的预测概率。
        y 是实际标签，通常是一个整数向量，表示每个样本的正确分类。
        self.length 保存了数据的样本数量，用于后续的计算。
        """
        self.y_hat = y_hat
        self.y = y
        self.length = len(y_hat)

    def cross_entropy(self):
        """
        交叉熵损失函数
        self.y_hat 是模型的预测输出，通常是经过 Softmax激活后的概率分布。
        它的形状为 (batch_size, num_classes)。
        self.y 是实际的标签，通常是一个一维数组，其中每个元素是对应样本的真实类别标签（即类别的索引）。
        self.length 是样本的数量。
        """
        y_tmp = self.y_hat[range(self.length), self.y]
        return -np.log(y_tmp)

    def cs_delta(self):
        """
        交叉熵+softmax求导
        yj是真实标签（one-hot 编码形式），它是一个0 或1的值
        对于 softmax + 交叉熵 的组合，损失函数 𝐿
        L 对输出层的第 j 个神经元的梯度（即导数）是：y_hat - y
        A_delta   softmax 输出误差的数组
        """
        A_delta = np.array(self.y_hat) #形状通常是 (batch_size, output_size)，其中 batch_size 是当前批次的样本数，10 是类别的数量（即输出层神经元的数量）
        A_delta[range(self.length), self.y] -= 1 #self.length 是当前批次样本的数量（batch_size）更新了交叉熵损失函数对输出概率的梯度
        #计算交叉熵损失函数对softmax输出的梯度。
        #选择A_delta中每个样本的真实标签对应的预测概率。
        #对于每个样本，将该位置的值减去1，表示交叉熵损失函数对该预测概率的梯度。
        return A_delta.mean(axis=0)#计算并返回梯度的均值

    def accuracy(self):
        """
        :return: 预测准确率
        """
        y_tmp = self.y_hat.argmax(axis=1)
        cmp = y_tmp.astype(dtype=self.y.dtype) == self.y
        cmp = cmp.astype(int).sum()
        return round(float(cmp) / len(self.y), 3)


class NeuralNetwork:
    """
    神经网络类
    """

    def __init__(self, layer_lst, training_data, test_data=None):
        # layer层数, 只算隐藏层和输出层
        self.nums_layer = len(layer_lst) - 1
        self.dims_layer = [(layer_lst[i + 1], layer_lst[i]) for i in range(self.nums_layer)]##元组中存储相邻两层神经元的个数
        # 训练样本个数
        self.num_training = 0
        self.W = []
        self.b = []
        # 存储前向传播梯度
        self.forward_grad = []
        # 存储后向传播梯度
        self.backward_grad = []

        self.training_features, self.training_labels = training_data[0], training_data[1]

        self.num_training = len(self.training_features)#训练样本个数

        if test_data:
            self.test_features, self.test_labels = test_data[0], test_data[1]
        else:
            self.test_features, self.test_labels = self.training_features, self.training_labels

        self.func1 = None
        self.func2 = None

    def set_activation(self, func1=Activation.ReLU, func2=Activation.Softmax):
        """
        设置激活函数
        :param func1: 隐藏层激活函数, 默认relu
        :param func2: 输出层激活函数, 默认softmax
        """
        self.func1 = func1
        self.func2 = func2

    def init_params(self):
        """
        初始化weight, bias, bp梯度
        """
        # 生成一个正态分布的随机数数组，均值为 0，标准差为 0.1
        self.W = [np.random.normal(0, 0.1, layer) for layer in self.dims_layer]
        self.b = [np.zeros((1, layer[0])) for layer in self.dims_layer]#为每一层生成一个零向量
        self.backward_grad = [np.zeros((layer[0], 1)) for layer in self.dims_layer]#反向传播的梯度初始化为零

    def data_iter(self, batch_size):
        """
        小批量梯度下降的迭代器
        :param batch_size: 批量大小, 在trainer方法传入
        """
        indices = list(range(self.num_training))#训练数据的索引
        random.shuffle(indices)

        for i in range(0, self.num_training, batch_size):
            idx = indices[i:min(i + batch_size, self.num_training)]#从indices列表中选择一段连续的子序列，表示当前批次的数据的索引。
            yield self.training_features[idx], self.training_labels[idx]#使用yield返回当前批次的特征和标签数据。

    def forward_propagation(self, X):
        """
        前向传播算法
        :param X: 小批量的features
        """
        self.forward_grad.append(X)

        for i in range(self.nums_layer):
            if i == self.nums_layer - 1:#输出层，使用softmax
                A = self.func2(self.forward_grad[i] @ self.W[i].T + self.b[i])#s（wx+b）
            else:#隐藏层，使用relu
                A = self.func1(self.forward_grad[i] @ self.W[i].T + self.b[i])#r(wx+b)

            self.forward_grad.append(A)

    def backward_propagation(self, y):
        """
        反向传播算法
        :param y: 小批量的labels
        """
        l = Loss(self.forward_grad[-1], y) #最后一层和预测值之间，计算loss

        self.backward_grad[-1] += l.cs_delta().reshape((-1, 1)) #交叉熵损失函数对输出层的导数（梯度）表示每个类别的误差（梯度）将一维数组重塑为一个 列向量

        d = Delta(self.func1)#激活函数relu的导数

        for i in reversed(range(self.nums_layer - 1)):
            # 将误差从输出层反向传播到每个隐藏层，并计算每层神经元的梯度，为权重更新做准备。
            tmp = d.delta(self.forward_grad[i + 1]).mean(axis=0) * (self.backward_grad[i + 1] * self.W[i + 1]).sum(
                axis=0)#进行计算上一层的误差

            tmp = tmp.reshape((-1, 1))#上一层的误差（梯度）

            self.backward_grad[i] += tmp#更新当前层的误差。

    def zero_grad(self):
        """
        梯度清零
        """
        self.forward_grad = []
        for i in range(self.nums_layer):
            self.backward_grad[i] = np.zeros_like(self.backward_grad[i])

    def SGD(self, lr):
        """
        梯度下降优化算法
        :param lr: 学习率, 在trainer方法传入
        w_grad：权重梯度（W的导数）。
        b_grad：偏置梯度（b的导数）。
        """
        for i in range(self.nums_layer):
            #梯度表示的是误差对于权重的导数，权重的更新应该由当前层的梯度信息和上一层的激活值决定。
            w_grad = self.forward_grad[i].mean(axis=0) * self.backward_grad[i]
            b_grad = self.backward_grad[i].reshape((1, -1))
            self.W[i] -= lr * w_grad#更新权重
            self.b[i] -= lr * b_grad#更新偏置

    def save_model(self, filename):
        """
        保存模型的参数（权重和偏置）以及训练历史（损失和准确率）
        :param filename: 保存文件的路径
        """
        model_data = {
            'W': self.W,  # 权重
            'b': self.b,  # 偏置
            'loss': self.loss_lst,  # 损失
            'accuracy': self.accuracy_lst  # 准确率
        }
        with open(filename, 'wb') as f:
            pickle.dump(model_data, f)
        print(f"模型已保存到 {filename}")

    def trainer(self, num_epochs, lr, batch_size):
        """
        训练器
        :param num_epochs: 训练的代数
        :param lr: 学习率
        :param batch_size: 批量大小
        """
        self.loss_lst, self.accuracy_lst = [], []

        self.init_params() #初始化权重、偏置、反向传播梯度

        for epoch in range(num_epochs):

            for X, y in self.data_iter(batch_size): #每次获取小批量的数据
                self.zero_grad()#梯度清零
                self.forward_propagation(X)#前向传播
                self.backward_propagation(y)#反向传播
                self.SGD(lr)#梯度下降，更新权重和偏置
                self.zero_grad()#梯度清零 在PyTorch框架中，通常每次反向传播后，梯度会自动清零

            self.forward_propagation(self.test_features)#对测试数据进行前向传播
            l = Loss(self.forward_grad[-1], self.test_labels)#计算loss

            loss = l.cross_entropy().mean()
            acc = l.accuracy()

            print("当前进行到第{}轮次，loss为{}，acc为{}".format(epoch, loss, acc))

            self.loss_lst.append(loss)
            self.accuracy_lst.append(acc)

        self.save_model('model/1/trained_model.pkl')  # 保存到文件

    def plt_show(self):
        """
        matplotlib画图
        """
        plt.figure(figsize=(14, 6))
        plt.subplot(1, 2, 1)
        plt.plot(self.loss_lst, linewidth=2.5, color='red', label='Loss')
        plt.title('Loss', fontsize=22, fontweight='bold', color='darkblue')
        plt.xlabel('Epochs', fontsize=16)
        plt.ylabel('Loss', fontsize=16)
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend(fontsize=14, loc='best')

        plt.subplot(1, 2, 2)
        plt.plot(self.accuracy_lst, linewidth=2.5, color='green', label='Accuracy')
        plt.title('Accuracy', fontsize=22, fontweight='bold', color='darkblue')
        plt.xlabel('Epochs', fontsize=16)
        plt.ylabel('Accuracy', fontsize=16)
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend(fontsize=14, loc='best')

        plt.tight_layout()
        plt.show()


def load_data(path):
    """
    加载MNIST数据集, 把每个图片展成784维的向量
    """
    f = gzip.open(path, 'rb')
    training_data, validation_data, test_data = _pickle.load(f, encoding='bytes')
    f.close()
    return [training_data, validation_data, test_data]


def load_model(filename):
    """
    加载保存的模型
    :param filename: 保存的模型文件路径
    :return: 加载的模型数据字典
    """
    with open(filename, 'rb') as f:
        model_data = pickle.load(f)
    print(f"模型已从 {filename} 加载")
    return model_data


def validate_model(model_data, test_data):
    """
    在测试集上验证模型
    :param model_data: 加载的模型数据（包括权重和偏置）
    :param test_data: 测试集数据（features 和 labels）
    """
    # 加载模型参数
    W = model_data['W']
    b = model_data['b']

    # 重新构建神经网络并设置参数
    nn = NeuralNetwork([784, 256, 256, 10], test_data, test_data)
    nn.W = W
    nn.b = b
    nn.func1 = Activation.ReLU
    nn.func2 = Activation.Softmax

    # 前向传播
    nn.forward_propagation(nn.test_features)
    l = Loss(nn.forward_grad[-1], nn.test_labels)
    loss = l.cross_entropy().mean()
    acc = l.accuracy()

    print(f"测试集上的损失：{loss}")
    print(f"测试集上的准确率：{acc}")

if __name__ == '__main__':
    # 加载数据集
    training_data, validation_data, test_data = load_data('MNIST/mnist.pkl.gz')

    # 超参数
    input_dims = 784
    hidden1_dims = 256
    hidden2_dims = 256
    output_dims = 10
    num_epochs = 200
    batch_size = 16
    lr = 0.02
    relu = Activation.ReLU
    softmax = Activation.Softmax

    # 构建神经网络并训练
    nn = NeuralNetwork([input_dims, hidden1_dims, hidden2_dims, output_dims], training_data, validation_data)#初始化神经网络的架构和数据
    nn.set_activation(relu, softmax) #设置激活函数隐藏层为relu，输出层为softmax
    nn.trainer(num_epochs, lr, batch_size)
    nn.plt_show()

    # 保存训练的模型
    nn.save_model('trained_model.pkl')

    # 加载保存的模型并验证
    model_data = load_model('model/1/trained_model.pkl')
    validate_model(model_data, test_data)
